#!/bin/python3
#此函数用于跟踪指定函数获取的参数
#用法：ffun "function(arg1,arg2，...)"
#需要指定函数声明，不需要指定函数类型，参数之间使用逗号分开

from bcc import BPF
from multiprocessing import cpu_count
import numpy as np
import sys
import os
import re
import sys
import signal


on_warn=False
create_struct=1

def printwarn(name='',msg=''):
    if not on_warn:
        return
    print(name+'\033[0;33;1m'+str(msg)+'\033[0m')
    #warnings.warn(str(msg),RuntimeWarning,stacklevel=2)

class arg_struct:
    def __init__(self):
        self.argstr=''  #参数字符串，格式“(xxx)”
        self.array=[]   #参数数组
        self.element_array=object
        self.name_array=object

#函数信息结构
class fun_struct:
    def __init__(self):
        self.str=''
        self.name=''
        self.id=0
        self.bpf_name=''
        self.bpf_str=''
        self.args=arg_struct()

class funa:
    funstr=''   #函数原型
    fun_name_index="funname"
    fun_args_index="funargs"
    fun_line=[]  #函数名和参数字典，由上面两个参数索引

    fun_reg=object  #函数拆分
    fun=object
    funs=[]     #函数数组
    funid=0     #函数id，从0开始

    bfun=object
    bfuns=[]    #bpf函数数组

    #bpf代码
    bpf_text=''
    bpf_out='events'    #信息输出管道名称

    #bpf函数组
    bpf_funs=[] #储存每一段函数
    bpf_ok=False    #代码可以使用

    mybpf=object

    #数据文件
    datafile_def_name='ffun.data'
    datafile_def_bakname='ffun.data.bak'
    datafile_name=datafile_def_name
    datafile_bakname=datafile_def_bakname
    #数据文件指针
    datafile_ptr=None

    #只有一个函数时，可以打印参数
    #enable_parameter=False
    #允许打印参数
    get_args=False

    #构造方法
    def __init__(self,cfun='',dfile=''):
        self.init_value()
        if cfun!='':
            self.set_fun(cfun)
            
    def init_file(self,datafile=''):    #初始化文件操作
        if datafile!='':
            self.datafile_name=datafile
            self.datafile_bakname=datafile+'.bak'
        #打开data文件，如果文件存在，增加后缀名”.bak“，如果不存在，重新打开
        try:
            os.rename(self.datafile_name,self.datafile_bakname)
        except Exception as err:
            printwarn(err)
        #现在文件不存在，可以创建
        self.datafile_ptr=open(self.datafile_name,'w+')
        return self.datafile_ptr

    def close(self):    #关闭输出文件
        self.datafile_ptr.close()
    
    def open(self,datafile=''):     #指定输出文件
        if self.datafile_ptr!=None:
            self.datafile_ptr.close()
        self.init_file(datafile)

    def init_value(self):   #初始化部分全局变量
        #创建拆分函数和参数模板
        self.fun_reg=re.compile(r'^(?P<'+self.fun_name_index+r'>[^ (]*)(\s*)(?P<'+self.fun_args_index+r'>[^;]*)')
    
    def init_args_text(self):
        self.bpf_text="""
        #include <linux/sched.h>
        BPF_PERF_OUTPUT("""+self.bpf_out+');'+"""
        struct data_t{
            u64 time;
            int pid;
            int tgid;
            int id;
            char comm[TASK_COMM_LEN];
        """
        for i in self.fun.args.element_array:
            if 'struct' in i:
                self.bpf_text+=' '.join(i)+';\n'
            else:
                self.bpf_text+=i[-2]+' '+i[-1]+';\n'
        self.bpf_text+='};\n'

    def init_text(self):
        #初始化bpf代码
        self.bpf_text="""
        #include <linux/sched.h>
        BPF_PERF_OUTPUT("""+self.bpf_out+');'+"""
        struct data_t{
            u64 time;
            int pid;
            int tgid;
            int id;
            char comm[TASK_COMM_LEN];
        };
        """

    def set_fun(self,funstr=''):    #初始化函数信息
        if funstr=='':
            return
        self.fun=fun_struct()
        self.fun.str=funstr
        #设置id
        self.fun.id=self.funid
        self.funid+=1
        if '(' and ')' not in funstr:    #参数部分有内容
            self.fun.name=funstr
            no_args=True
        else:
        #拆分函数
            self.fun_line=self.fun_reg.match(funstr).groupdict()
            self.fun.name=self.fun_line[self.fun_name_index]
            no_args=False
        #bpf函数名
        #self.bfun=fun_struct()
        #self.bfun.name='my_bpf_'+self.fun.name
        self.fun.bpf_name='my_bpf_'+self.fun.name
        #bpf函数声明
        self.fun.bpf_str='int '+self.fun.bpf_name
        while True:
            if no_args:
                self.fun.bpf_str+='(struct pt_regs *ctx)'
                break
            #拆分参数
            try:
                self.fun.args.array=re.split(r'\s*[,\(\)]\s*',self.fun_line[self.   fun_args_index])[1:-1]
            except Exception as err:
                #if self.fun.args.array==['']:
                self.fun.bpf_str+='(struct pt_regs *ctx)'
                break
            args_len=len(self.fun.args.array)
            self.fun.args.element_array=np.empty([args_len],dtype=object)
            self.fun.args.name_array=np.empty([args_len],dtype=(str,255))
            for i in range(args_len):
                self.fun.args.element_array[i]=re.sub(r'\s*?\*','* ',self.fun.args.array[i]).split()
                self.fun.args.name_array[i]=self.fun.args.element_array[i][-1]
            self.fun.bpf_str+=self.fun_line[self.fun_args_index].replace('(','(struct   pt_regs *ctx, ')
            #将相关结构添加到数组
            
            #self.bfuns.append(self.bfun)
            break
        self.funs.append(self.fun)
        #生成bpf函数段
        if self.get_args:
            self.create_bfun(self.fun.bpf_str)
        else:
            self.create_def_bfun(self.fun.bpf_str)

    bfun_def_text1="""
        {
            struct data_t data={};
            u64 ptid=bpf_get_current_pid_tgid();
            data.pid=ptid;
            data.tgid=ptid>>32;
            bpf_get_current_comm(&data.comm,sizeof(data.comm));
            data.time=bpf_ktime_get_ns();
            data.id="""

    def create_bfun(self,bfun_str):  #创建bpf函数
        tmp=bfun_str+self.bfun_def_text1+str(self.fun.id)+';\n'
        try:
            for i in self.fun.args.element_array:
                tmp+='data.'+i[-1]+'='+i[-1]+';\n'
            tmp+=self.bpf_out+'.perf_submit(ctx,&data,sizeof(data));'+"""
                return 0;
            }
            """
            self.bpf_funs.append(tmp)
        except Exception as err:
            printwarn("没有参数")
    
    def create_def_bfun(self,bfun_str):
        tmp=bfun_str+self.bfun_def_text1+str(self.fun.id)+';\n'+self.bpf_out+'.perf_submit(ctx,&data,sizeof(data));'+"""
            return 0;
        }
        """
        self.bpf_funs.append(tmp)

    def text(self):     #生成bpf代码块
        if self.get_args:
            self.init_args_text()
        else:
            self.init_text()
        for i in self.bpf_funs:
            self.bpf_text+=i
        self.bpf_ok=True

    def runbpf(self):   #执行bpf代码
        if not self.bpf_ok:
            self.text()
        self.mybpf=BPF(text=self.bpf_text)
        return self.mybpf

    def aa_kprobe(self):    #挂到指定函数上
        for i in range(len(self.funs)):
            self.bpf.attach_kprobe(event=self.funs[i].name,fn_name=self.funs[i].bpf_name)
            #self.bpf.attach_kfunc(fn_name=self.funs[i].bpf_name)
            #self.bpf.attach_kretprobe(event=self.funs[i].name,fn_name=self.funs[i].bpf_name)

    def ns2s(self,ns): #将纳秒转换为秒
        return float(ns)/1000000000

    last_time=0

    def print_file(self,all=[]):
        self.datafile_ptr.writelines('%d\t%f\t%d\t\n' % (all[0],all[1],all[2]))

    def print_data(self,cpu,data,size):
        event=self.bpf[self.bpf_out].event(data)
        new_time=self.ns2s(event.time)
        #由于阻塞或者目标函数调用过于频繁，会发生当前时间的数据先输出，前一条才输出的情况，导致差值可能出现负值，读取数据时要注意
        subnl=new_time-self.last_time
        self.print_file([cpu,subnl,event.id])
        self.last_time=new_time

    def print_data_args(self,cpu,data,size):
        event=self.bpf[self.bpf_out].event(data)
        new_time=self.ns2s(event.time)
        subnl=new_time-self.last_time
        self.datafile_ptr.writelines('%d\t%f\t%s' % (cpu,subnl,event.id))
        for i in self.fun.args.element_array:
            self.datafile_ptr.writelines('\t'+str(eval('event.'+i[-1])))
        self.datafile_ptr.write('\n')
        self.last_time=new_time

    #中断函数
    def sigint_handler(self,signum,frame):
        self.bpf.cleanup()
        print()
        signal.alarm(0)
        self.datafile_ptr.flush()
        self.datafile_ptr.close()
        
        os._exit(0)

    #@tn:运行时间“秒”,默认运行5s
    def run(self,tn=5):
        self.bpf=self.runbpf()
        self.aa_kprobe()    #挂上探测代码
        self.datafile_ptr.writelines('cpus=%d\n'% cpu_count()) #写入cpu数量
        self.datafile_ptr.writelines('funs=%d\n'% len(self.funs))   #写入函数数量
        for i in self.funs: #写入函数名及其编号
            self.datafile_ptr.writelines('%s\t%d\n'%(i.name,i.id))
        #signal.signal(signal.SIGINT,self.sigint_handler)
        signal.signal(signal.SIGALRM,self.sigint_handler)
        signal.alarm(tn)
        if self.get_args:
            self.bpf[self.bpf_out].open_perf_buffer(self.print_data_args)
        else:
            self.bpf[self.bpf_out].open_perf_buffer(self.print_data)
        while True:
            try:
                self.bpf.perf_buffer_poll()
            except KeyboardInterrupt:
                return
        
    def addfun(self,sss=''):
        if sss=='':
            return ''
        self.set_fun(sss)
        return sss

import getopt

def main(argv):

    if len(argv) == 0 :
        return
    time_num=5  #运行时间
    function_sss=[]     #函数数组
    print_bpf=False   #是否打印bpf函数
    get_args=False    #输出函数参数
    not_run=False       #禁止执行bpf代码

    main_options='ht:f:pvn'
    main_options_name=['help','time=','function=','print','not-run']

    try:
        opts,args=getopt.getopt(argv,main_options,main_options_name)
    except getopt.GetoptError as err:
        sys.exit()
    for opt,arg in opts:
        if opt in ('-h','--help'):
            print('python3 ffun.py -f <function>')
            print('-t\t指定运行时间，单位\'秒\'')
            return
        elif opt in ('-t','--time'):
            time_num=int(arg)
        elif opt in ('-f','--function'):
            function_sss.append(arg)
        elif opt in ('-p','--print'):
            print_bpf=True
        elif opt in ('-v'):
            get_args=True
        elif opt in ('-n','--not-run'):
            not_run=True

    f1=funa()
    if  len(function_sss)==1:
        if '(' and ')' in function_sss[0]:
            f1.get_args=get_args        #允许输出参数
        f1.addfun(function_sss[0])
    else:
        if get_args:
            print('函数超过一个，不能使用\'-v\'选项')
        for i in function_sss:
            f1.addfun(i)

    #f1.addfun('find_busiest_group(u64 env)')
    f1.text()
    if print_bpf:
        print(f1.bpf_text)
    f1.open()
    if not not_run:
        f1.runbpf()
        f1.run(time_num)

if __name__=='__main__':
    main(sys.argv[1:])